import os

import scipy.misc

import cv2 as cv
import copy
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt

if __name__ == '__main__':
    L_path = r"output/predictions_LR/" # 机器预测的灰度图所在路径
    raw_path = r""  # 原始图片所在路径
    aff_path = r""  # 人工分割图片所在路径
    save_path = r""  # 输出保存位置
    png_names_list = os.listdir(raw_path)  # 目录下灰度图片的列表
    LR_names_list = os.listdir(L_path)


    for i in range(0, len(png_names_list)):
        png_dir = raw_path + png_names_list[i]
        raw_image = cv.imread(png_dir, 0)
        png_dir = aff_path + png_names_list[i]
        aff_image = cv.imread(png_dir, 0)
        png_dir = L_path + LR_names_list[i]
        prediction_image = cv.imread(png_dir, 0)
        a1 = copy.deepcopy(prediction_image)
        a2 = copy.deepcopy(prediction_image)
        a3 = copy.deepcopy(prediction_image)
        a1[a1 == 1] = 255
        a1[a1 == 2] = 0
        a2[a2 == 1] = 255
        a2[a2 == 2] = 255
        a3[a3 == 1] = 255
        a3[a3 == 2] = 77
        a1 = Image.fromarray(np.uint8(a1)).convert('L')
        a2 = Image.fromarray(np.uint8(a2)).convert('L')
        a3 = Image.fromarray(np.uint8(a3)).convert('L')
        prediction = Image.merge('RGB', [a1, a2, a3])
        output_name = os.path.splitext(png_names_list[i])[0]
        prediction.save(save_path + output_name + '.png')
        prediction_f = np.array(prediction)
        final_image_s = np.concatenate((raw_image, aff_image, prediction_f), axis=1)
        final_image = Image.fromarray(final_image_s)
        final_image.save(save_path + "/concact/" + output_name + '_concact.png')

